import argparse
import json
import os,utils
import tqdm
from datetime import datetime
import logging
import asyncio
import argparse
import json
import os
import asyncio
import aiohttp
import tqdm.asyncio as tqdm_async
from typing import Dict, Any


# 禁用所有不必要的日志
logging.getLogger("openai").setLevel(logging.WARNING)
logging.getLogger("httpx").setLevel(logging.WARNING)
logging.getLogger("httpcore").setLevel(logging.WARNING)
logging.getLogger("urllib3").setLevel(logging.WARNING)


example = '''example 1:
    question:
    static int ossl_statem_server13_read_transition(SSL_CONNECTION *s, int mt)
{
    OSSL_STATEM *st = &s->statem;

    /*
     * Note: There is no case for TLS_ST_BEFORE because at that stage we have
     * not negotiated TLSv1.3 yet, so that case is handled by
     * ossl_statem_server_read_transition()
     */

    switch (st->hand_state) {
    default:
        break;

    case TLS_ST_EARLY_DATA:
        if (s->hello_retry_request == SSL_HRR_PENDING) {
            if (mt == SSL3_MT_CLIENT_HELLO) {
                st->hand_state = TLS_ST_SR_CLNT_HELLO;
                return 1;
            }

    case TLS_ST_CR_CERT_REQ:
        if (s->post_handshake_auth == SSL_PHA_REQUESTED) {
            if (do_compressed_cert(s))
                st->hand_state = TLS_ST_CW_COMP_CERT;
            else
                st->hand_state = TLS_ST_CW_CERT;
            return WRITE_TRAN_CONTINUE;
        
    Reasoning process:
    In the switch statement, the second case checks if TLS_ST_EARLY_DATA is one of the values in the given state list. Therefore, the initial state is TLS_ST_EARLY_DATA. Within this case, there are two nested if statements: if (s->hello_retry_request == SSL_HRR_PENDING) and if (mt == SSL3_MT_CLIENT_HELLO). Consequently, the condition is "s->hello_retry_request == SSL_HRR_PENDING AND mt == SSL3_MT_CLIENT_HELLO." Additionally, since st->hand_state = TLS_ST_SR_CLNT_HELLO, the final state is TLS_ST_SR_CLNT_HELLO.

    answer:

{
    "function": "ossl_statem_server13_read_transition",
    "transitions": [
        {
            "start_state": "TLS_ST_EARLY_DATA",
            "condition": "s->hello_retry_request == SSL_HRR_PENDING and mt == SSL3_MT_CLIENT_HELLO",
            "condition explanation": "server need client send Hello Retry Request __AND__ server read message is SSL3_MT_CLIENT_HELLO"
            "end_state": "TLS_ST_SR_CLNT_HELLO"
        },
        {
            "start_state": "TLS_ST_CR_CERT_REQ",
            "condition": "s->post_handshake_auth == SSL_PHA_REQUESTED and do_compressed_cert(s)",
            "condition explanation":"post handshake authentication was requested __AND__ compressed certificate was requested",
            "end_state": "TLS_ST_CW_COMP_CERT"
        },
        {
            "start_state": "TLS_ST_CR_CERT_REQ",
            "condition": "s->post_handshake_auth == SSL_PHA_REQUESTED and not do_compressed_cert(s)",
            "condition explanation":"post handshake authentication was requested __AND__ __NOT__ compressed certificate was requested",
            "end_state": "TLS_ST_CW_CERT"
        }
    ]
}


            '''



get_fsm_prompt = lambda x,y: f'You are an experienced protocol expert. The protocol implementation is {x} Provide a function to extract state transitions from it, ensuring that all transitions are from one state to another, rather than a list of states. This point is crucial. The output format should be: {{"start_state":start state,"condition":condition of transition,"condition explanation":condition explanation in English,"end_state":end state}}.  The logical connective in condition explanation should be replaced with the following dictionary {{"and":"__AND__", "or":"__OR__","not":"__NOT__"}}  think step by step. {example}, \n\n question : {y}'  




def setup_logger(log_path):
    """Set up logging configuration"""
    # 创建 logger 对象（而不是直接用 basicConfig）
    logger = logging.getLogger()
    logger.setLevel(logging.DEBUG)  # 全局最低级别设为 DEBUG
    
    # 清除默认的 handlers（避免重复日志）
    logger.handlers.clear()
    
    # 1. 文件处理器（记录 DEBUG 及以上级别）
    file_handler = logging.FileHandler(log_path, encoding="utf-8")
    file_handler.setLevel(logging.DEBUG)  # 文件记录所有 DEBUG+ 日志
    file_formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
    file_handler.setFormatter(file_formatter)
    
    # 2. 控制台处理器（仅记录 INFO 及以上级别）
    console_handler = logging.StreamHandler()
    console_handler.setLevel(logging.INFO)  # 终端只打印 INFO+
    console_formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
    console_handler.setFormatter(console_formatter)
    
    # 添加处理器
    logger.addHandler(file_handler)
    # logger.addHandler(console_handler)
    
    return logger
            


async def process_function(group_name: str, func_name: str, func_data: Dict, 
                          protocol_impl: str, logger, api_key: str, semaphore):
    """处理单个函数的异步任务"""
    async with semaphore:  # 控制并发数
        try:
            prompt = get_fsm_prompt(protocol_impl, func_data)
            logger.info(f"Sending to LLM: {func_name}")
            
            answer = await utils.get_ans_json_async(prompt)
            
            logger.debug(f"Response for {func_name}:\n{json.dumps(answer, indent=2)}")
            return (group_name, func_name, answer)
            
        except Exception as e:
            logger.error(f"Error in {func_name}: {str(e)}")
            return (group_name, func_name, {"error": str(e)})


async def async_main(config):
    """异步主函数"""
    save_dir = os.path.join(os.path.abspath(config['save_parent_dir']), 
                          f"{config['project']}_{config['version']}")
    os.makedirs(save_dir, exist_ok=True)
    
    # 初始化日志
    log_path = os.path.join(save_dir, 'get_fsm_from_llm.log')
    logger = setup_logger(log_path)
    logger.info(f"Starting async processing for {config['project']}")
    
    # 加载配置
    with open(config['state_functions_path'], encoding='utf-8') as f:
        state_funcs_config = json.load(f)
    
    # 准备任务
    semaphore = asyncio.Semaphore(10)  # 最大并发数
    tasks = []
    for group_name, func_names in config['state_functions'].items():
        for func_name in func_names:
            task = process_function(
                group_name, func_name, 
                state_funcs_config[func_name],
                config['project'], logger,
                config['api_key'], semaphore
            )
            tasks.append(task)
    
    # 并行执行（带进度条）
    results = []
    for f in tqdm_async.tqdm.as_completed(tasks, desc="Processing"):
        result = await f
        results.append(result)
    
    # 整理结果
    llm_fsm_response = {}
    for group_name, func_name, answer in results:
        if group_name not in llm_fsm_response:
            llm_fsm_response[group_name] = {}
        llm_fsm_response[group_name][func_name] = answer
    
    # 保存结果
    save_path = os.path.join(save_dir, "get_fsm_from_llm.json")
    with open(save_path, 'w', encoding='utf-8') as f:
        json.dump(llm_fsm_response, f, indent=2)
    
    logger.info(f"Saved results to {save_path}")

# python get_fsm_from_llm.py ./config/openssl.json
if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Async code analysis')
    parser.add_argument('config', help='Path to config JSON')
    args = parser.parse_args()
    
    with open(args.config, encoding='utf-8') as f:
        config = json.load(f)
    
    asyncio.run(async_main(config))